import os
import torch
from torch.utils import data
from models import build_model
import matplotlib.pyplot as plt
import torch.nn.functional as F
from MyDatasets import kaggle
import torchvision.transforms as transforms
import yaml
import easydict as Edict
from PIL import Image

label_dict = {0 : "cat",
              1 : "dog"}

def test_net(configs):
    model = build_model(configs.backbone, num_classes=configs.Num_Classes, pretrained=False)
    with open("./weights/final.pth", 'rb') as f:
        state_dict = torch.load(f)
    model.load_state_dict(state_dict)
    model.eval()
    if configs.cuda:
        device = torch.device("cuda")
        model.to(device)
    if configs.img_aug:
        imgaug = transforms.Compose([
                                     transforms.CenterCrop(224),
                                     transforms.ToTensor()])
        test_set = kaggle(configs.test_root, transform=imgaug)
        test_loader = data.DataLoader(test_set, batch_size=configs.Test.batch_size,
                                                   shuffle=configs.shuffle, num_workers=configs.num_workers, pin_memory=True)
    else:
        test_set = kaggle(configs.test_root, transform=None)
        test_loader = data.Dataloader(test_set, batch_size=configs.Test.batch_size,
                                                   shuffle=configs.shuffle, num_workers=configs.num_workers,
                                                   pin_memory=True)

    for idx, (names, img) in enumerate(test_loader):
        if configs.cuda:
            device = torch.device("cuda")
            img = img.to(device)

        out = model(img)
        out = torch.argmax(F.softmax(out), dim=1).cpu().numpy()   # labels
        for i in range(configs.Test.batch_size):
            img = Image.open(os.path.join(configs.test_root, names[i]))
            plt.subplot(4, 4, i + 1)
            plt.imshow(img)
            plt.text(50, 50, label_dict[out[i]])
            plt.title(names[i])
        plt.show()


    with open("./weights/final.pth", 'wb') as f:
        torch.save(model.state_dict(), f)
        print("Final model saved!!!")

if __name__ == "__main__":
    with open("./config.yaml") as f:
        configs = yaml.load(f)
        configs = Edict.EasyDict(configs)
    test_net(configs)
